package com.example.springai.commonConfig;

import com.google.gson.Gson;
import com.google.gson.JsonElement;
import com.google.gson.JsonObject;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.*;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;

import java.util.Collections;
import java.util.List;


@Component
@RequiredArgsConstructor
public class RedisChatMemory implements ChatMemory {
    private static final String REDIS_CHAT_MEMORY_KEY = "chatmemory";
    private final StringRedisTemplate redisTemplate;
    private static final Gson GSON = new Gson();

    @Override
    public void add(String conversationId, List<Message> messages) {
        String key = REDIS_CHAT_MEMORY_KEY + ":" + conversationId;
        List<String> serializedMessages = messages.stream().map(this::serialize).toList();
        if (redisTemplate.hasKey(key)){
            redisTemplate.opsForList().rightPushAll(key, serializedMessages);
        } else {
            redisTemplate.opsForList().rightPushAll(key, serializedMessages);
            redisTemplate.expire(key, 7 * 24 * 60 * 60, java.util.concurrent.TimeUnit.SECONDS);
        }
    }

    @Override
    public List<Message> get(String conversationId) {
        String key = REDIS_CHAT_MEMORY_KEY + ":" + conversationId;
        List<String> serializedMessages = redisTemplate.opsForList().range(key, 0, -1); // 获取全部
        if (CollectionUtils.isEmpty(serializedMessages)) {
            return Collections.emptyList();
        }
        return serializedMessages.stream().map(this::deserialize).toList();
    }

    public List<Message> getLastN(String conversationId, int lastN) {
        String key = REDIS_CHAT_MEMORY_KEY + ":" + conversationId;
        List<String> serializedMessages = redisTemplate.opsForList().range(key, -lastN, -1);
        if (CollectionUtils.isEmpty(serializedMessages)) {
            return Collections.emptyList();
        }
        return serializedMessages.stream().map(this::deserialize).toList();
    }

    @Override
    public void clear(String conversationId) {
        redisTemplate.delete(REDIS_CHAT_MEMORY_KEY + ":" + conversationId);
    }

    private String serialize(Message message) {
        return GSON.toJson(message);
    }

    private Message deserialize(String serializedMessage) {
        JsonObject jsonObject = GSON.fromJson(serializedMessage, JsonObject.class);
        JsonElement typeElement = jsonObject.get("messageType");
        String messageType = typeElement.getAsString().trim(); // 去除前后空格和引号干扰
        return switch (messageType) {
            case "USER" -> GSON.fromJson(serializedMessage, UserMessage.class);
            case "ASSISTANT" -> GSON.fromJson(serializedMessage, AssistantMessage.class);
            case "SYSTEM" -> GSON.fromJson(serializedMessage, SystemMessage.class);
            case "TOOL" -> GSON.fromJson(serializedMessage, ToolResponseMessage.class);
            default -> throw new IllegalStateException("Unexpected value: " + messageType);
        };
    }
}
